package cas.ihep.spark.io;

import java.io.ByteArrayOutputStream;
import java.io.DataInput;
import java.io.EOFException;
import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.nio.ByteBuffer;

@SuppressWarnings("WeakerAccess")
public final class ByteArrayOutputStreamExt extends ByteArrayOutputStream implements DataInput{

    private int readOffset;
    private boolean isClosed;
    private static final Method ensureCapacity;

    static{
        try{
            ensureCapacity=ByteArrayOutputStream.class.getDeclaredMethod("ensureCapacity", int.class);
            ensureCapacity.setAccessible(true);
        }catch(NoSuchMethodException e){
            throw new RuntimeException(e);
        }
    }
    public ByteArrayOutputStreamExt(int size){
        super(size);
        isClosed=false;
    }

    public void close(){
        if(!isClosed){
            isClosed=true;
        }
    }

    public void write(int b){
        if(isClosed) {
            throw new RuntimeException("The stream is closed");
        }
        super.write(b);
    }

    public synchronized void write(byte[] buf,int off,int len){
        if(isClosed) {
            throw new RuntimeException("The stream is closed");
        }
        super.write(buf,off,len);
    }

    public byte[] toByteArray(){
        if(isClosed) {
            return buf;
        }
        throw new RuntimeException("Can only get array buffer after the stream is closed");
    }

    public void write(ByteBuffer buffer){
        if(isClosed) {
            throw new RuntimeException("The stream is closed");
        }
        int remaining=buffer.remaining();
        int beforePos=buffer.position();
        try{
            ensureCapacity.invoke(this,count+remaining);
        }catch (IllegalAccessException | InvocationTargetException e){
            throw new RuntimeException(e);
        }
        buffer.get(buf,count,remaining);
        buffer.position(beforePos);
        count+=remaining;
    }

    @Override
    public void readFully(byte[] b) throws IOException {
        readFully(b,0,b.length);
    }

    @Override
    public void readFully(byte[] b, int off, int len) throws IOException {
        int r=pread(readOffset,b,off,len);
        if(r<len){
            throw new EOFException();
        }
        readOffset+=r;
    }

    public int read(byte[] b,int off,int len){
        int r=pread(readOffset,b,off,len);
        if(r!=-1) {
            readOffset += r;
        }
        return r;
    }

    public int pread(int rOff,byte[] b,int off,int len){
        if(b==null){
            throw new NullPointerException();
        }
        if(rOff<0 || rOff>=count || off<0 || off+len>b.length){
            throw new IndexOutOfBoundsException();
        }
        len=Math.min(count-rOff,len);
        System.arraycopy(buf,rOff,b,off,len);
        return len;
    }

    public int pread(int rOff,ByteBuffer dst){
        if(rOff<0 || rOff>count){
            throw new IndexOutOfBoundsException();
        }
        int avail=Math.min(dst.remaining(),count-rOff);
        dst.put(buf,rOff,avail);
        return avail;
    }

    public int read(ByteBuffer dst){
        int r=pread(readOffset,dst);
        if(r>0){
            readOffset+=r;
        }
        return r;
    }

    public int readAndRemove(byte[] b){
        return readAndRemove(b,0,b.length);
    }

    public int readAndRemove(ByteBuffer b){
        int r=pread(readOffset,b);
        if(r>0){
            int off=readOffset+r;
            for(int len=readOffset;off<count;len++,off++){
                buf[len]=buf[off];
            }
            count-=r;
        }
        return r;
    }

    public int readAndRemove(byte[] b,int off,int len){
        int r=pread(readOffset,b,off,len);
        if(r>0){
            off=readOffset+r;
            for(len=readOffset;off<count;len++,off++){
                buf[len]=buf[off];
            }
            count-=r;
        }
        return r;
    }

    public int pread(int off){
        return off<count?buf[off]:-1;
    }

    public int read(){
        return pread(readOffset++);
    }

    @Override
    public int skipBytes(int n) {
        if(n<0){
            n=Math.min(readOffset,-n);
            readOffset-=n;
        }else {
            n = Math.min(n, count - readOffset);
            readOffset += n;
        }
        return n;
    }

    @Override
    public boolean readBoolean() throws IOException {
        int r=read();
        if(r!=-1){
            return r!=0;
        }
        throw new EOFException();
    }

    public boolean preadBoolean(int off)throws IOException{
        int r=pread(off);
        if(r!=-1){
            return r!=0;
        }
        throw new EOFException();
    }

    @Override
    public byte readByte() throws IOException {
        int r=read();
        if(r!=-1){
            return (byte)r;
        }
        throw new EOFException();
    }

    public byte preadByte(int off)throws IOException{
        int r=pread(off);
        if(r!=-1){
            return (byte)r;
        }
        throw new EOFException();
    }

    @Override
    public int readUnsignedByte() throws IOException {
        return readByte()&0xff;
    }

    @Override
    public short readShort() throws IOException {
        if(readOffset+2<=count){
            short ret=byteArray2Short(buf,readOffset);
            readOffset+=2;
            return ret;
        }
        throw new EOFException();
    }

    @Override
    public int readUnsignedShort() throws IOException {
        return readShort()&0xffff;
    }

    @Override
    public char readChar() throws IOException {
        return (char)readUnsignedShort();
    }

    @Override
    public int readInt() throws IOException {
        if(readOffset+4<=count){
            int ret=byteArray2Int(buf,readOffset);
            readOffset+=4;
            return ret;
        }
        throw new EOFException();
    }

    public int preadInt(int rOff)throws IOException{
        if(rOff<0){
            throw new IndexOutOfBoundsException();
        }
        if(rOff+4<=count){
            return byteArray2Int(buf,rOff);
        }
        throw new EOFException();
    }

    @Override
    public long readLong() throws IOException {
        if (readOffset+8<=count){
            long ret=byteArray2Long(buf,readOffset);
            readOffset+=8;
            return ret;
        }
        throw new EOFException();
    }

    @Override
    public float readFloat() throws IOException {
        return Float.intBitsToFloat(readInt());
    }

    @Override
    public double readDouble() throws IOException {
        return Double.longBitsToDouble(readLong());
    }

    @Override
    public String readLine() {
        throw new RuntimeException("Unsupported operation");
    }

    @Override
    public String readUTF() {
        throw new RuntimeException("Unsupported operation");
    }

    private static void int2ByteArray(int v,byte[] buf,int offset){
        buf[offset]= (byte) ((v>>>24)&0xff);
        buf[offset+1]= (byte) ((v>>>16)&0xff);
        buf[offset+2]= (byte) ((v>>>8)&0xff);
        buf[offset+3]= (byte) (v&0xff);
    }

    private static void short2ByteArray(short sh,byte[] buf,int offset){
        buf[offset]= (byte) ((sh>>>8)&0xff);
        buf[offset+1]= (byte) (sh&0xff);
    }

    private static void long2ByteArray(long v,byte[] buf,int offset){
        buf[offset]= (byte) ((v>>>56)&0xff);
        buf[offset+1]= (byte) ((v>>>48)&0xff);
        buf[offset+2]= (byte) ((v>>>40)&0xff);
        buf[offset+3]= (byte) ((v>>>32)&0xff);
        buf[offset+4]= (byte) ((v>>>24)&0xff);
        buf[offset+5]= (byte) ((v>>>16)&0xff);
        buf[offset+6]= (byte) ((v>>>8)&0xff);
        buf[offset+7]= (byte) (v&0xff);
    }
    private static short byteArray2Short(byte arr[],int off){
        return (short) (arr[off]<<8 | arr[off+1]);
    }

    private static int byteArray2Int(byte[] arr,int off){
        return arr[off]<<24 | arr[off+1]<<16 | arr[off+2]<<8 | arr[off+3];
    }

    private static long byteArray2Long(byte arr[],int off){
        return (long)arr[off]<<56 | (long)arr[off+1]<<48 | (long)arr[off+2]<<40 | (long)arr[off+3]<<32 |
                arr[off+4]<<24 | arr[off+5]<<16 | arr[off+6]<<8 | arr[off+7];
    }
}
